import torch
from torch import nn

IMG_SIZE = 224
LR = 0.9

class VGG(nn.Module):
    """VGG Base"""
    def __init__(self, in_channels, conv_arch, num_classes):
        super(VGG, self).__init__()

        blocks = []
        for num_convs, out_channels in conv_arch:
            blocks.append(self.vgg_block(num_convs, in_channels, out_channels))
            in_channels = out_channels
        
        self.net = nn.Sequential(
            *blocks,
            nn.Flatten(),
            nn.Linear(out_channels * 7 * 7, 4096), nn.Dropout(),
            nn.Linear(4096, 4096), nn.Dropout(),
            nn.Linear(4096, 1000), nn.Dropout(),
            nn.Linear(1000, num_classes) 
        )


    def vgg_block(self, num_convs, in_channels, out_channels):
        seq = nn.Sequential()
        for _ in range(num_convs):
            seq.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1))
            seq.append(nn.ReLU())
            in_channels = out_channels
        seq.append(nn.MaxPool2d(kernel_size=2, stride=2))
        return seq


    def forward(self, input):
        return self.net(input)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR
    

conv_arch = {
    11: [(1, 64), (1, 128), (2, 256), (2, 512), (2, 512)],
    13: [(2, 64), (2, 128), (2, 256), (2, 512), (2, 512)],
    16: [(2, 64), (2, 128), (3, 256), (3, 512), (3, 512)],
    19: [(2, 64), (2, 128), (3, 256), (4, 512), (4, 512)]
}


class VGG11(VGG):
    """VGG11"""
    def __init__(self, num_classes, in_channels):
        super(VGG11, self).__init__(in_channels, conv_arch[11], num_classes)


class VGG13(VGG):
    """VGG13"""
    def __init__(self, num_classes, in_channels):
        super(VGG13, self).__init__(in_channels, conv_arch[13], num_classes)


class VGG16(VGG):
    """VGG16"""
    def __init__(self, num_classes, in_channels):
        super(VGG16, self).__init__(in_channels, conv_arch[16], num_classes)


class VGG16C(nn.Module):
    """VGG16 C with conv1"""
    def __init__(self, num_classes, in_channels):
        super(VGG16C, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(512 * 8 * 8, 4096), nn.Dropout(),
            nn.Linear(4096, 4096), nn.Dropout(),
            nn.Linear(4096, 1000), nn.Dropout(),
            nn.Linear(1000, num_classes) 
        )


    def forward(self, input):
        return self.net(input)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class VGG19(VGG):
    """VGG19"""
    def __init__(self, num_classes, in_channels):
        super(VGG19, self).__init__(in_channels, conv_arch[19], num_classes)